<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>LMQL Chat</title>
    <link rel="stylesheet" href="chat_assets/index.css">
    <script>
    function addMessage(message, type, id=null) {
        if (message.trim() === '') {
            return;
        }

        const chat = document.querySelector('.chat');

        let wasThinking = isThinking;
        setThinking(false)

        // check if message already exists, then update it
        if (id && document.querySelector(`.message#${id}`)) {
            let existingMessage = document.querySelector(`.message#${id}`)
            // content 
            let content = existingMessage.querySelector('.message-content');
            // check if scrolled to bottom
            let isScrolledToBottom = chat.scrollHeight - chat.clientHeight <= chat.scrollTop + 1;

            content.innerHTML = message;
            
            if (isScrolledToBottom) {
                chat.scrollTop = chat.scrollHeight;
            }
            return;
        }

        const messageElement = document.createElement('div');
        messageElement.classList.add('message');
        messageElement.classList.add(type);
        messageElement.id = id;
        const messageContent = document.createElement('div');
        messageContent.classList.add('message-content');
        messageContent.innerText = message;
        messageElement.appendChild(messageContent);
        chat.appendChild(messageElement);
        chat.scrollTop = chat.scrollHeight;
        if (wasThinking) {
            setThinking(wasThinking);
        }
    }
    function clickSend() {
        const input = document.querySelector('.input textarea');
        const message = input.value;
        input.value = '';
        addMessage(message, 'assistant');
    }
    function chatKeyDown() {
        const input = document.querySelector('.input textarea');
        if (event.keyCode === 13 && !event.shiftKey) {
            event.preventDefault();
            sendMessage();
        }
    }
    let isThinking = false;
    function setThinking(state) {
        if (!isThinking && state) {
            isThinking = true;
            addThinking();
        } else if (isThinking && !state) {
            isThinking = false;
            removeThinking();
        }
    }
    function addThinking() {
        // add jumping bubbles to chat
        const chat = document.querySelector('.chat');
        const thinking = document.createElement('div');
        thinking.classList.add('thinking');
        // add three bubbles
        for (let i = 0; i < 3; i++) {
            const bubble = document.createElement('div');
            bubble.classList.add('bubble');
            thinking.appendChild(bubble);
        }
        chat.appendChild(thinking);
    }
    function removeThinking() {
        // remove jumping bubbles from chat
        const chat = document.querySelector('.chat');
        const thinking = document.querySelector('.thinking');
        chat.removeChild(thinking);
    }
    function sendMessage() {
        if (isThinking) {
            return;
        }
        const input = document.querySelector('.input textarea');
        if (input.value.trim() === '') {
            return;
        }
        const message = input.value;
        input.value = '';
        addMessage(message, 'user');
        setThinking(true);
        document.querySelector('.input textarea').disabled = true;
        send(message);
    }

    let messageId = 0;
    let websocket = null;
    

    window.onload = function() {
        setThinking(true);
        
        document.getElementById('connection-status').innerText = 'Connecting...';
        
        host = window.location.host;
        websocket = new WebSocket("ws://" + host + "/chat");
        document.getElementById('address').innerText = websocket.url;
        
        websocket.onmessage = function(event) {
            let payload = JSON.parse(event.data);
            if (payload.type === 'response') {
                let id = payload.message_id;
                let object = payload.data;
                // show prompt on backside
                let prompt = object.prompt;
                // replace all <lmql:*/> tags with <span class="lmql-tag">...</span>
                prompt = prompt.replace(/<lmql:.*?\/>/g, (match) => {
                    // replace angled brackets with html entities
                    let tag_type = match.substring(6, match.length - 2);
                    return `<span class="lmql-tag ${tag_type}">${tag_type}</span>`;
                });
                // escape all << and >>
                prompt = prompt.replace(/<</g, '&lt;&lt;');
                prompt = prompt.replace(/>>/g, '&gt;&gt;');
                document.getElementById("raw-prompt").innerHTML = prompt;

                addResponse(object.variables["ANSWER"].trim(), 'assistant', "response-" + id);
            }
        }
        websocket.onopen = function(event) {
            document.getElementById('connection-status').innerText = 'Connected';
            websocket.send(JSON.stringify({
                type: 'init'
            }));
            setThinking(false);
        }
        websocket.onclose = function(event) {
            document.getElementById('connection-status').innerText = 'Disconnected';
        }

        // get send and focus
        document.querySelector('.input textarea').focus();
    }

    function send(question) {
        setThinking(true);
        let id = messageId++;

        websocket.send(JSON.stringify({
            type: 'input',
            text: question
        }));
    }

    function addResponse(s, role, id=null) {
        document.querySelector('.input textarea').disabled = false;
        addMessage(s, role, id);
        document.querySelector('.input textarea').focus();
    }
    </script>
</head>
<body>
    <div class="content side-view">
        <div class="front side">
        <div class="toolbar">
            <img src="chat_assets/lmql.svg" />
            <h1>LMQL Chat</h1>
            <h2 id="connection-status"></h2>
            <button onclick="document.querySelector('.content').classList.toggle('side-view')">
                <img src="chat_assets/code.svg"/>
            </button>
        </div>
        <div class="chat"></div>
        <div class="input">
            <textarea onkeydown="chatKeyDown()"></textarea>
            <button onclick="clickSend()">
                <img src="chat_assets/send.svg"/>
            </button>
        </div>
        </div>
        <div class="back code side">
            <div class="toolbar">
                <h1>Internal Trace</h1>
                <h2 id="address"></h2>
                <!-- <button onclick="document.querySelector('.content').classList.toggle('side-view')">
                    <img src="chat_assets/code.svg"/>
                </button> -->
            </div>
            <div id="raw-prompt">
            </div>
        </div>
    </div>
</body>
</html>